package com.zytech.zspider.transformer;

import com.google.common.base.Joiner;
import com.google.common.base.Predicate;
import com.google.common.collect.*;
import com.zytech.modules.lang.Langs;
import com.zytech.modules.lang.Strings;
import com.zytech.zspider.core.*;
import org.antlr.v4.runtime.tree.TerminalNode;

import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * @author siuming
 */
class SpiderDefinitionTransformerListener extends SpiderDefinitionBaseListener {

    private static final String FATHER_COL = "$father";
    private static final String TASKDEF_COL = "$taskdef";
    private static final String CHILDREN_COL = "$children";

    private SpiderDefinition def;
    private SpiderDefinition.Builder defBuilder;

    private String curtaskName;
    private String curtaskUrl;
    private String curtaskMethod;
    private TaskDefinition curtask;
    private Table<String, String, Object> tasks;

    private TaskGeneratorTuple curtaskTuple;
    private List<TaskGeneratorTuple> generatorTuples;
    private Map<String, Boolean> generatorTupleRefs = Maps.newHashMap();

    /**
     * @return
     */
    public SpiderDefinition getDef() {
        return def;
    }

    @Override
    public void enterTransform(SpiderDefinitionParser.TransformContext ctx) {
        defBuilder = SpiderDefinition.ofBuilder();
        tasks = HashBasedTable.create();
        generatorTuples = Lists.newArrayList();
    }

    @Override
    public void exitTransform(SpiderDefinitionParser.TransformContext ctx) {

        Map<String, Map<String, Object>> rows = tasks.rowMap();
        Map<String, Map<String, Object>> filter = Maps.filterEntries(rows, new Predicate<Map.Entry<String, Map<String, Object>>>() {
            @Override
            public boolean apply(Map.Entry<String, Map<String, Object>> input) {
                return null == input.getValue().get(FATHER_COL);
            }
        });

        if (filter.size() != 1) {
            throw new BadSpiderDefinitionException("more than one root TaskDefinition.");
        }

        Map.Entry<String, Map<String, Object>> entry = filter.entrySet().iterator().next();
        String taskname = entry.getKey();
        TaskDefinition task = (TaskDefinition) entry.getValue().get(TASKDEF_COL);
        constructTasks(taskname, task);
        def = defBuilder.task(task).build();
    }

    private void constructTasks(String taskname, TaskDefinition task) {
        Set<String> subtaskNames = (Set<String>) tasks.row(taskname).get(CHILDREN_COL);
        if (subtaskNames.isEmpty()) {
            return;
        }

        for (String subtaskName : subtaskNames) {
            Object subtask = tasks.row(subtaskName).get(TASKDEF_COL);
            if (null == subtask) {
                throw new TaskDefinitionNotFoundException(subtaskName);
            }
            task.subtask((TaskDefinition) subtask);
            constructTasks(subtaskName, (TaskDefinition) subtask);
        }
    }

    @Override
    public void enterNameDef(SpiderDefinitionParser.NameDefContext ctx) {
        if (ctx.getParent() instanceof SpiderDefinitionParser.TransformContext) {
            defBuilder.name(trimText(ctx.STRING().getText()));
            return;
        }

        if (ctx.getParent() instanceof SpiderDefinitionParser.SubtasksGeneratorDefContext) {
            curtaskTuple.name = trimText(ctx.STRING().getText());
            return;
        }

        curtaskName = trimText(ctx.STRING().getText());
    }

    private String trimText(String text) {
        String trims = text;
        if (Strings.startsWith(trims, "'''")) {
            trims = trims.substring(3);
        } else if (Strings.startsWith(trims, "'") || Strings.startsWith(trims, "\"")) {
            trims = trims.substring(1);
        }

        if (Strings.endsWith(trims, "'''")) {
            trims = trims.substring(0, trims.length() - 3);
        } else if (Strings.endsWith(trims, "'") || Strings.endsWith(trims, "\"")) {
            trims = trims.substring(0, trims.length() - 1);
        }
        return trims;
    }

    @Override
    public void enterUrlDef(SpiderDefinitionParser.UrlDefContext ctx) {
        curtaskUrl = trimText(ctx.STRING().getText());
    }

    @Override
    public void enterMethodDef(SpiderDefinitionParser.MethodDefContext ctx) {
        curtaskMethod = trimText(ctx.STRING().getText());
    }

    @Override
    public void exitMethodDef(SpiderDefinitionParser.MethodDefContext ctx) {
        curtask = TaskDefinition.of(curtaskName, curtaskUrl, curtaskMethod);
        if (tasks.row(curtaskName).containsKey(TASKDEF_COL)) {
            throw new TaskDefinitionExistedException(curtaskName);
        }

        tasks.put(curtaskName, TASKDEF_COL, curtask);
        tasks.put(curtaskName, CHILDREN_COL, Sets.newHashSet());
    }

    @Override
    public void enterHeadersDef(SpiderDefinitionParser.HeadersDefContext ctx) {
        if (ctx.getParent() instanceof SpiderDefinitionParser.TransformContext) {
            Map<String, String> headers = Maps.newHashMap();
            for (SpiderDefinitionParser.PairContext pairContext : ctx.pair()) {
                String name = getPairName(pairContext);
                String value = getPairValue(pairContext);
                headers.put(name, value);
            }

            defBuilder.headers(headers);
            return;
        }

        for (SpiderDefinitionParser.PairContext pairContext : ctx.pair()) {
            String name = getPairName(pairContext);
            String value = getPairValue(pairContext);
            curtask.header(name, value);
        }
    }

    private String getPairName(SpiderDefinitionParser.PairContext pairContext) {
        return trimText(pairContext.STRING(0).getText());
    }

    private String getPairValue(SpiderDefinitionParser.PairContext pairContext) {
        if (null != pairContext.INT()) {
            return pairContext.INT().getText();
        }
        return trimText(pairContext.STRING(1).getText());
    }

    @Override
    public void enterParamsDef(SpiderDefinitionParser.ParamsDefContext ctx) {
        if (ctx.getParent() instanceof SpiderDefinitionParser.TransformContext) {
            Map<String, String> params = Maps.newHashMap();
            for (SpiderDefinitionParser.PairContext pairContext : ctx.pair()) {
                String name = getPairName(pairContext);
                String value = getPairValue(pairContext);
                params.put(name, value);
            }

            defBuilder.params(params);
            return;
        }

        for (SpiderDefinitionParser.PairContext pairContext : ctx.pair()) {
            String name = getPairName(pairContext);
            String value = getPairValue(pairContext);
            curtask.param(name, value);
        }
    }

    @Override
    public void enterParamsGeneratorDef(SpiderDefinitionParser.ParamsGeneratorDefContext ctx) {
        Map<String, String> paramsGenerator = Maps.newHashMap();
        for (SpiderDefinitionParser.PairContext pairContext : ctx.pair()) {
            String name = getPairName(pairContext);
            String value = getPairValue(pairContext);
            paramsGenerator.put(name, value);
        }

        defBuilder.paramsGenerator(paramsGenerator);
    }

    @Override
    public void enterBodyDef(SpiderDefinitionParser.BodyDefContext ctx) {
        curtask.body(trimText(ctx.STRING().getText()));
    }

    @Override
    public void enterBodyGeneratorDef(SpiderDefinitionParser.BodyGeneratorDefContext ctx) {
        defBuilder.bodyGenerator(trimText(ctx.STRING().getText()));
    }

    @Override
    public void enterSettingsDef(SpiderDefinitionParser.SettingsDefContext ctx) {
        if (ctx.getParent() instanceof SpiderDefinitionParser.TransformContext) {
            Map<String, String> props = Maps.newHashMap();
            for (SpiderDefinitionParser.PairContext pairContext : ctx.pair()) {
                String name = getPairName(pairContext);
                String value = getPairValue(pairContext);
                props.put(name, value);
            }

            defBuilder.settings(props);
            return;
        }

        for (SpiderDefinitionParser.PairContext pairContext : ctx.pair()) {
            String name = getPairName(pairContext);
            String value = getPairValue(pairContext);
            curtask.setting(name, value);
        }
    }

    @Override
    public void enterOnstartDef(SpiderDefinitionParser.OnstartDefContext ctx) {
        if (ctx.getParent() instanceof SpiderDefinitionParser.SubtasksGeneratorDefContext) {
            curtaskTuple.startScript = trimText(ctx.STRING().getText());
            return;
        }
        curtask.onStart(trimText(ctx.STRING().getText()));
    }

    @Override
    public void enterOnresultDef(SpiderDefinitionParser.OnresultDefContext ctx) {
        TaskResultPipelineSettings.Builder builder = TaskResultPipelineSettings.ofBuilder();
        builder.process(trimText(ctx.processDef().STRING().getText()));
        if (null != ctx.retryDef()) {
            builder.retryWhen(trimText(ctx.retryDef().STRING().getText())).retryTimes(3);
            if (null != ctx.retryTimesDef()) {
                builder.retryTimes(Integer.valueOf(ctx.retryTimesDef().INT().getText()));
            }
        }

        if (ctx.getParent() instanceof SpiderDefinitionParser.SubtasksGeneratorDefContext) {
            curtaskTuple.resultHandler = builder.build();
            return;
        }

        curtask.onResult(builder.build());
    }

    @Override
    public void enterOnfinishDef(SpiderDefinitionParser.OnfinishDefContext ctx) {
        if (ctx.getParent() instanceof SpiderDefinitionParser.SubtasksGeneratorDefContext) {
            curtaskTuple.finishScript = trimText(ctx.STRING().getText());
            return;
        }
        curtask.onFinish(trimText(ctx.STRING().getText()));
    }

    @Override
    public void enterNextDef(SpiderDefinitionParser.NextDefContext ctx) {
        for (TerminalNode node : ctx.STRING()) {
            String next = trimText(node.getText());
            curtaskTuple.nexts.add(next);
            if (!generatorTupleRefs.containsKey(next)) {
                generatorTupleRefs.put(next, Boolean.TRUE);
            }
        }
    }

    @Override
    public void enterSubtasksGeneratorDef(SpiderDefinitionParser.SubtasksGeneratorDefContext ctx) {
        curtaskTuple = new TaskGeneratorTuple();
    }

    @Override
    public void exitSubtasksGeneratorDef(SpiderDefinitionParser.SubtasksGeneratorDefContext ctx) {
        curtaskTuple.script = trimText(ctx.STRING().getText());
        generatorTuples.add(curtaskTuple);
    }

    @Override
    public void enterSubtasksDef(SpiderDefinitionParser.SubtasksDefContext ctx) {
        Set<String> subtaskNames = (Set<String>) tasks.row(curtaskName).get(CHILDREN_COL);

        for (TerminalNode node : ctx.STRING()) {
            String subname = trimText(node.getText());
            tasks.put(subname, FATHER_COL, curtaskName);
            subtaskNames.add(subname);
        }
    }

    @Override
    public void enterCollectsDef(SpiderDefinitionParser.CollectsDefContext ctx) {
        if (ctx.getParent() instanceof SpiderDefinitionParser.SubtasksGeneratorDefContext) {
            for (SpiderDefinitionParser.PairContext pairContext : ctx.pair()) {
                String name = getPairName(pairContext);
                String script = getPairValue(pairContext);
                curtaskTuple.collectScripts.put(name, script);
            }

            return;
        }
        for (SpiderDefinitionParser.PairContext pairContext : ctx.pair()) {
            String name = getPairName(pairContext);
            String script = getPairValue(pairContext);
            curtask.collect(name, script);
        }
    }

    @Override
    public void exitTaskDef(SpiderDefinitionParser.TaskDefContext ctx) {
        if (generatorTuples.isEmpty()) {
            return;
        }

        Iterable<TaskGeneratorTuple> noRef = Iterables.filter(generatorTuples, new Predicate<TaskGeneratorTuple>() {
            @Override
            public boolean apply(TaskGeneratorTuple input) {
                return !generatorTupleRefs.containsKey(input.name);
            }
        });

        for (TaskGeneratorTuple tuple : noRef) {
            TaskGenDefinition generator = constructTaskGenerator(tuple);
            curtask.subtaskGenDefinition(generator);
        }

        curtaskTuple = null;
        generatorTuples.clear();
    }

    private TaskGenDefinition constructTaskGenerator(final TaskGeneratorTuple tuple) {
        List<TaskGeneratorTuple> nextTuples = Langs.filter(generatorTuples, new com.zytech.modules.lang.Predicate<TaskGeneratorTuple>() {
            @Override
            public boolean apply(TaskGeneratorTuple input) {
                return tuple.nexts.contains(input.name);
            }
        });

        if (nextTuples.size() != tuple.nexts.size()) {
            throw new BadTaskDefinitionException(
                    String.format(
                            "some TaskGenDefinition[%s] cann't be found.",
                            Joiner.on(',').join(tuple.nexts)
                    )
            );
        }

        TaskGenDefinition.Builder builder = TaskGenDefinition
                .ofBuilder()
                .name(tuple.name)
                .onStart(tuple.startScript)
                .onResult(tuple.resultHandler)
                .onFinish(tuple.finishScript)
                .onCollects(tuple.collectScripts)
                .onGenerate(tuple.script);

        for (TaskGeneratorTuple nextTuple : nextTuples) {
            TaskGenDefinition next = constructTaskGenerator(nextTuple);
            builder.next(next);
        }

        return builder.build();
    }

    class TaskGeneratorTuple {
        String name;
        String script;
        String startScript;
        String finishScript;
        List<String> nexts = Lists.newArrayList();
        Map<String, String> collectScripts = Maps.newHashMap();
        TaskResultPipelineSettings resultHandler;
    }
}
